{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4d8329ce-c442-4226-bc0f-183097bed6f6",
   "metadata": {},
   "source": [
    "## Procedural memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "40bcb974-0d02-49ec-9eb7-cf2a820c3ac8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture cap --no-stderr\n",
    "%pip install -U langmem langgraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "131c57c5-570e-4bb0-a9d4-0aa00ea89ccd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "store = InMemoryStore()\n",
    "store.put((\"instructions\",), key=\"agent_instructions\", value={\"prompt\": \"Write good emails.\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c2e7817c-f876-4f70-80ce-1131f74c8b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "from langgraph.config import get_store\n",
    "\n",
    "def draft_email(to: str, subject: str, body: str):\n",
    "    \"\"\"Submit an email draft.\"\"\"\n",
    "    return \"Draft saved successfully.\"\n",
    "\n",
    "def prompt(state):\n",
    "    item = store.get((\"instructions\",), key=\"agent_instructions\")\n",
    "    instructions = item.value[\"prompt\"]\n",
    "    sys_prompt = {\"role\": \"system\", \"content\": f\"## Instructions\\n\\n{instructions}\"}\n",
    "    return [sys_prompt] + state['messages']\n",
    "\n",
    "agent = create_react_agent(\"anthropic:claude-3-5-sonnet-latest\", prompt=prompt, tools=[draft_email], store=store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e9b3a50b-93c0-4ac8-9de0-4f4058d872fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "\n",
      "[{'citations': None, 'text': \"I'll help you draft an email to schedule the follow-up meeting.\", 'type': 'text'}, {'id': 'toolu_01Reo41SH7Sx72LVzt2iCEBi', 'input': {'to': 'joe@langchain.dev', 'subject': 'Follow-up Meeting - Thursday at Noon', 'body': \"Hi Joe,\\n\\nI'd like to schedule a follow-up meeting with you for Thursday at 12:00 PM. \\n\\nPlease let me know if this time works for you.\\n\\nBest regards\"}, 'name': 'draft_email', 'type': 'tool_use'}]\n",
      "Tool Calls:\n",
      "  draft_email (toolu_01Reo41SH7Sx72LVzt2iCEBi)\n",
      " Call ID: toolu_01Reo41SH7Sx72LVzt2iCEBi\n",
      "  Args:\n",
      "    to: joe@langchain.dev\n",
      "    subject: Follow-up Meeting - Thursday at Noon\n",
      "    body: Hi Joe,\n",
      "\n",
      "I'd like to schedule a follow-up meeting with you for Thursday at 12:00 PM. \n",
      "\n",
      "Please let me know if this time works for you.\n",
      "\n",
      "Best regards\n"
     ]
    }
   ],
   "source": [
    "result = agent.invoke(\n",
    "    {\"messages\": [\n",
    "        {\"role\": \"user\", \"content\" :\"Draft an email to joe@langchain.dev saying that we want to schedule a followup meeting for thursday at noon.\"}]}\n",
    ")\n",
    "result['messages'][1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5159932a-23a1-421f-9119-07b9f4dc5e69",
   "metadata": {},
   "source": [
    "# Update the prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2aa4675d-afff-45f7-8836-da0c132b8caa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langmem import create_prompt_optimizer\n",
    "\n",
    "optimizer = create_prompt_optimizer(\"anthropic:claude-3-5-sonnet-latest\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b80f55a9-f8f1-4b01-adeb-f24152b5c465",
   "metadata": {},
   "outputs": [],
   "source": [
    "current_prompt = store.get((\"instructions\",), key=\"agent_instructions\").value[\"prompt\"]\n",
    "feedback = {\"request\": \"Always sign off from 'William'; for meeting requests, offer to schedule on Zoom or Google Meet\"}\n",
    "\n",
    "optimizer_result = optimizer.invoke({\"prompt\": current_prompt, \"trajectories\": [(result[\"messages\"], feedback)]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "357c3086-bad6-46cc-81cf-9b2442940c92",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Write professional emails following these guidelines:\n",
      "1. Always sign off with 'William'\n",
      "2. For meeting requests:\n",
      "   - Include the meeting time and date\n",
      "   - Offer both Zoom and Google Meet as platform options\n",
      "3. Keep the tone professional but friendly\n",
      "4. Include a clear subject line\n",
      "\n",
      "Example meeting request:\n",
      "Subject: Follow-up Meeting - Thursday at Noon\n",
      "Body:\n",
      "Hi [Name],\n",
      "\n",
      "I'd like to schedule a follow-up meeting with you for [Day] at [Time].\n",
      "\n",
      "We can meet via Zoom or Google Meet - please let me know your preference.\n",
      "\n",
      "Please let me know if this time works for you.\n",
      "\n",
      "Best regards,\n",
      "William\n"
     ]
    }
   ],
   "source": [
    "print(optimizer_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "07c09ab8-4649-4603-a760-55a6e77559f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "store.put((\"instructions\",), key=\"agent_instructions\", value={\"prompt\": optimizer_result})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "38b190a0-a72a-4894-a2ad-eeee9657edc5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "\n",
      "[{'citations': None, 'text': \"I'll help you draft a meeting request email to Joe following the guidelines.\", 'type': 'text'}, {'id': 'toolu_01BUKzWhkWje1tfyJ4BjRVdo', 'input': {'to': 'joe@langchain.dev', 'subject': 'Follow-up Meeting - Thursday at Noon', 'body': \"Hi Joe,\\n\\nI'd like to schedule a follow-up meeting with you for Thursday at 12:00 PM.\\n\\nWe can meet via Zoom or Google Meet - please let me know your preference.\\n\\nPlease let me know if this time works for you.\\n\\nBest regards,\\nWilliam\"}, 'name': 'draft_email', 'type': 'tool_use'}]\n",
      "Tool Calls:\n",
      "  draft_email (toolu_01BUKzWhkWje1tfyJ4BjRVdo)\n",
      " Call ID: toolu_01BUKzWhkWje1tfyJ4BjRVdo\n",
      "  Args:\n",
      "    to: joe@langchain.dev\n",
      "    subject: Follow-up Meeting - Thursday at Noon\n",
      "    body: Hi Joe,\n",
      "\n",
      "I'd like to schedule a follow-up meeting with you for Thursday at 12:00 PM.\n",
      "\n",
      "We can meet via Zoom or Google Meet - please let me know your preference.\n",
      "\n",
      "Please let me know if this time works for you.\n",
      "\n",
      "Best regards,\n",
      "William\n"
     ]
    }
   ],
   "source": [
    "result = agent.invoke(\n",
    "    {\"messages\": [\n",
    "        {\"role\": \"user\", \"content\" :\"Draft an email to joe@langchain.dev saying that we want to schedule a followup meeting for thursday at noon.\"}]}\n",
    ")\n",
    "result['messages'][1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d3cbfc93-5d1f-407d-b0a9-4f86d875ceb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "\n",
      "[{'citations': None, 'text': \"I'll help you draft an email to Roger about the release timing.\", 'type': 'text'}, {'id': 'toolu_01RHJf4D6T96t9VVo5xwDQGy', 'input': {'to': 'roger@langchain.dev', 'subject': 'Release Update - 4:00 PM Today', 'body': 'Hi Roger,\\n\\nI wanted to let you know that the release should be completed by 4:00 PM today.\\n\\nPlease let me know if you have any questions.\\n\\nBest regards,\\nWilliam'}, 'name': 'draft_email', 'type': 'tool_use'}]\n",
      "Tool Calls:\n",
      "  draft_email (toolu_01RHJf4D6T96t9VVo5xwDQGy)\n",
      " Call ID: toolu_01RHJf4D6T96t9VVo5xwDQGy\n",
      "  Args:\n",
      "    to: roger@langchain.dev\n",
      "    subject: Release Update - 4:00 PM Today\n",
      "    body: Hi Roger,\n",
      "\n",
      "I wanted to let you know that the release should be completed by 4:00 PM today.\n",
      "\n",
      "Please let me know if you have any questions.\n",
      "\n",
      "Best regards,\n",
      "William\n"
     ]
    }
   ],
   "source": [
    "result = agent.invoke(\n",
    "    {\"messages\": [\n",
    "        {\"role\": \"user\", \"content\" : \"Let roger@langchain.dev know that the release should be later by 4:00 PM.\"}]}\n",
    ")\n",
    "result['messages'][1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5630d1f7-9772-4471-bca0-22558d4978e9",
   "metadata": {},
   "source": [
    "## Multi-agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "52a85066-00eb-4986-a894-389b07d78da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture cap --no-stderr\n",
    "%pip install -U langgraph-supervisor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7d7a8b09-fe7d-4053-9a73-6a7d0e227404",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "from langgraph.config import get_store\n",
    "\n",
    "store = InMemoryStore()\n",
    "\n",
    "store.put((\"instructions\",), key=\"email_agent\", value={\"prompt\": \"Write good emails. Repeat your draft content to the user after submitting.\"})\n",
    "store.put((\"instructions\",), key=\"twitter_agent\", value={\"prompt\": \"Write fire tweets. Repeat the tweet content to the user upon submission.\"})\n",
    "\n",
    "## Email agent\n",
    "def draft_email(to: str, subject: str, body: str):\n",
    "    \"\"\"Submit an email draft.\"\"\"\n",
    "    return \"Draft saved successfully.\"\n",
    "\n",
    "def prompt_email(state):\n",
    "    item = store.get((\"instructions\",), key=\"email_agent\")\n",
    "    instructions = item.value[\"prompt\"]\n",
    "    sys_prompt = {\"role\": \"system\", \"content\": f\"## Instructions\\n\\n{instructions}\"}\n",
    "    return [sys_prompt] + state['messages']\n",
    "\n",
    "email_agent = create_react_agent(\n",
    "    \"anthropic:claude-3-5-sonnet-latest\", \n",
    "    prompt=prompt_email, \n",
    "    tools=[draft_email], \n",
    "    store=store,\n",
    "    name=\"email_assistant\",\n",
    ")\n",
    "\n",
    "## Tweet\n",
    "\n",
    "def tweet(to: str, subject: str, body: str):\n",
    "    \"\"\"Post a tweet.\"\"\"\n",
    "    return \"Legendary.\"\n",
    "\n",
    "def prompt_social_media(state):\n",
    "    item = store.get((\"instructions\",), key=\"twitter_agent\")\n",
    "    instructions = item.value[\"prompt\"]\n",
    "    sys_prompt = {\"role\": \"system\", \"content\": f\"## Instructions\\n\\n{instructions}\"}\n",
    "    return [sys_prompt] + state['messages']\n",
    "\n",
    "social_media_agent = create_react_agent(\n",
    "    \"anthropic:claude-3-5-sonnet-latest\", \n",
    "    prompt=prompt_social_media, \n",
    "    tools=[tweet], \n",
    "    store=store,\n",
    "    name=\"social_media_agent\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "76f062c4-2f5f-4a9a-b4dc-a23ec4651b36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph_supervisor import create_supervisor\n",
    "\n",
    "# Create supervisor workflow\n",
    "workflow = create_supervisor(\n",
    "    [email_agent, social_media_agent],\n",
    "    model=\"anthropic:claude-3-5-sonnet-latest\",\n",
    "    prompt=(\n",
    "        \"You are a team supervisor managing email and tweet assistants to help with correspondence.\"\n",
    "    )\n",
    ")\n",
    "\n",
    "# Compile and run\n",
    "app = workflow.compile(store=store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ff861c5a-d9ff-4905-b902-6b5550801539",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = app.invoke(\n",
    "    {\"messages\": [\n",
    "        {\"role\": \"user\", \"content\" :\"Draft an email to joe@langchain.dev saying that we want to schedule a followup meeting for thursday at noon.\"}]},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "16857af2-23bf-4a1c-a55b-7ae205decdc5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "Name: email_assistant\n",
      "\n",
      "I've drafted the email with the following content:\n",
      "\n",
      "To: joe@langchain.dev\n",
      "Subject: Follow-up Meeting - Thursday at Noon\n",
      "Body:\n",
      "Hi Joe,\n",
      "\n",
      "I'd like to schedule a follow-up meeting with you for Thursday at noon. Please let me know if this time works for you.\n",
      "\n",
      "Looking forward to our discussion.\n",
      "\n",
      "Best regards\n",
      "\n",
      "The email has been drafted and is ready to be sent. Let me know if you'd like any changes to the content.\n"
     ]
    }
   ],
   "source": [
    "result[\"messages\"][3].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d29809be-0c4e-47b8-befe-1c5ef71594bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langmem import create_multi_prompt_optimizer\n",
    "\n",
    "feedback = {\"request\": \"Always sign off emails from 'William'; for meeting requests, offer to schedule on Zoom or Google Meet\"}\n",
    "\n",
    "optimizer = create_multi_prompt_optimizer(\"anthropic:claude-3-5-sonnet-latest\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6734520b-f0be-428f-b538-c17824d2ec25",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001B[0;31mInit signature:\u001B[0m \u001B[0mPrompt\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m/\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
       "\u001B[0;31mDocstring:\u001B[0m     \n",
       "TypedDict for structured prompt management and optimization.\n",
       "\n",
       "Example:\n",
       "    ```python\n",
       "    from langmem import Prompt\n",
       "\n",
       "    prompt = Prompt(\n",
       "        name=\"extract_entities\",\n",
       "        prompt=\"Extract key entities from the text:\",\n",
       "        update_instructions=\"Make minimal changes, only address where\"\n",
       "        \" errors have occurred after reasoning over why they occur.\",\n",
       "        when_to_update=\"If there seem to be errors in recall of named entities.\",\n",
       "    )\n",
       "    ```\n",
       "\n",
       "The name and prompt fields are required. Optional fields control optimization:\n",
       "- update_instructions: Guidelines for modifying the prompt\n",
       "- when_to_update: Dependencies between prompts during optimization\n",
       "\n",
       "Use in the prompt optimizers.\n",
       "\u001B[0;31mFile:\u001B[0m           ~/code/lc/community/langmem/src/langmem/prompts/types.py\n",
       "\u001B[0;31mType:\u001B[0m           _TypedDictMeta\n",
       "\u001B[0;31mSubclasses:\u001B[0m     "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from langmem import Prompt\n",
    "?Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1813fa94-501b-43a0-8d20-c44ee2c037ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "email_prompt = store.get((\"instructions\",), key=\"email_agent\").value['prompt']\n",
    "tweet_prompt = store.get((\"instructions\",), key=\"twitter_agent\").value['prompt']\n",
    "\n",
    "email_prompt = {\n",
    "    \"name\": \"email_prompt\",\n",
    "    \"prompt\": email_prompt,\n",
    "    \"when_to_update\": \"Only if feedback is provided indicating email writing performance needs improved.\"\n",
    "}\n",
    "tweet_prompt = {\n",
    "    \"name\": \"tweet_prompt\",\n",
    "    \"prompt\": tweet_prompt,\n",
    "    \"when_to_update\": \"Only if tweet writing generation needs improvement.\"\n",
    "}\n",
    "\n",
    "\n",
    "optimizer_result = optimizer.invoke({\"prompts\": [tweet_prompt, email_prompt], \"trajectories\": [(result[\"messages\"], feedback)]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "bbc9d37a-23de-4d46-a7c4-219ba84b7625",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'name': 'tweet_prompt',\n",
       "  'prompt': 'Write fire tweets. Repeat the tweet content to the user upon submission.',\n",
       "  'when_to_update': 'Only if tweet writing generation needs improvement.'},\n",
       " {'name': 'email_prompt',\n",
       "  'prompt': \"Write professional emails following these requirements:\\n1. Include a clear subject line and appropriate greeting\\n2. Write concise and clear content\\n3. For meeting requests, always include Zoom or Google Meet as platform options\\n4. Always sign off emails with 'Best regards, William'\\n5. Repeat your draft content to the user after submitting\",\n",
       "  'when_to_update': 'Only if feedback is provided indicating email writing performance needs improved.'}]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9ef3d1fd-00b4-42a6-b56b-cd9b2d135b8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "store.put((\"instructions\",), key=\"email_agent\", value={\"prompt\": optimizer_result[1]['prompt']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d8d66705-a569-43cc-8c45-668cdec62f9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = app.invoke(\n",
    "    {\"messages\": [\n",
    "        {\"role\": \"user\", \"content\" :\"Draft an email to joe@langchain.dev saying that we want to schedule a followup meeting for thursday at noon.\"}]},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "11ae5952-01ce-4374-bd38-c9a22ec44d8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "Name: email_assistant\n",
      "\n",
      "I've drafted the email with:\n",
      "\n",
      "To: joe@langchain.dev\n",
      "Subject: Follow-up Meeting Request - Thursday at Noon\n",
      "Body:\n",
      "Hi Joe,\n",
      "\n",
      "I hope this email finds you well. I would like to schedule a follow-up meeting with you for this Thursday at 12:00 PM.\n",
      "\n",
      "We can meet via Zoom or Google Meet, whichever you prefer. Please let me know which platform works best for you, and I'll send over the meeting details.\n",
      "\n",
      "Looking forward to our discussion.\n",
      "\n",
      "Best regards,\n",
      "William\n",
      "\n",
      "The email has been drafted with a clear subject line, professional greeting, meeting time details, video conferencing options, and appropriate sign-off. Would you like me to modify anything in the draft?\n"
     ]
    }
   ],
   "source": [
    "result[\"messages\"][3].pretty_print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
